#include "yolov3_gpu_detector.h"
#include "imagebind_inference.h"
#include <dirent.h>
#include <opencv2/dnn/dnn.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/core/cuda.hpp>
#include <opencv2/opencv.hpp>
#include <fstream>
#include <sstream>
#include <cstdio>

void ReadImage(std::string imgdirs, std::vector<std::string>&image_paths){
  DIR *dir;
  if ((dir = opendir(imgdirs.c_str())) == nullptr){
    throw std::runtime_error("dircetory " + imgdirs + "does not exit");
  }
  dirent * dp;
  for (dp = readdir(dir); dp !=nullptr; dp=readdir(dir)){
    const std::string file_name = dp->d_name;
    if (file_name == "." || file_name==".."){
      continue;
    }
    image_paths.push_back(file_name);
  }
  closedir(dir);
}

int main(int argc, char ** argv)
{
  bool flag;
  std::string imgdir = "/home/admincaode/code/imagebind_yolo/main/data";
  std::string modelpath = "/home/admincaode/code/imagebind_yolo/main/weights";
  std::string imagebind_model = "";
  yoloV3GPUDetector yolo = yoloV3GPUDetector(modelpath);
  imagebindInfer imagebind = imagebindInfer();
  flag = imagebind.Initialize(imagebind_model);
  if (!flag){
    std::cout<<"init imagebind flase"<<std::endl;
    throw "init flase";
  }
  std::vector<std::string> imgfiles;
  ReadImage(imgdir, imgfiles);
  for (auto name: imgfiles){
    std::vector<cv::Rect2d> vecRoi;
    std::vector<std::string> vecName;
    std::vector<float> confidenc;
    std::string image_path = imgdir +"/"+ name;
    std::cout<<" "<<std::endl;
    std::cout<<image_path<<std::endl;
    cv::Mat cv_img = cv::imread(image_path, cv::IMREAD_COLOR);
    cv::Mat new_img = cv_img.clone();
    flag = imagebind.inference(cv_img,yolo, vecRoi, vecName, confidenc);
    if(flag == 0){
      printf( "~ TEST-1: Detect nothing lalalalalala~\n");
    }else{
      int bb_num = vecName.size();
      // DEBUG
      for (int i=0; i < confidenc.size(); i++){
        printf("~~~~~~~~ Detection Class Num = %d, Name = %s, score = %f ~~~~~~~~\n", bb_num, vecName[i].c_str(), confidenc[i]);
      }
      for (int i=0; i < vecRoi.size(); i++)
      {
        cv::rectangle(new_img, vecRoi[i], cv::Scalar(0, 0, 255), 2, 1);
        std::string s = vecName[i] + "_" + std::to_string(confidenc[i]);
        cv::Point origin;
        origin.x = vecRoi[i].x -5;
        origin.y = vecRoi[i].y -5;
        cv::putText(new_img, s, origin, cv::FONT_HERSHEY_COMPLEX,1, cv::Scalar(255, 0, 0), 1, 8, 0);
      }
      std::string savepath = "/home/admincaode/code/imagebind_yolo/main/outputs/" + name;
      //std::string savepath = "/home/yk/xjl_learn/test_data_0603/528/0423_del.jpg";
      printf("savepath: %s ", savepath.c_str());
      cv::imwrite(savepath,new_img);
      printf("******************************************************************************************");
    }
  }
}